import numpy as np
import pandas as pd
import torch
import time
import torch.nn as nn  # 这里将重要的库命名。
import torch.nn.functional as F  # 这里将内置函数也命名。
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader  # 这里引入一个dataloader.


# 在这里定义一个类。
class CNN_net(nn.Module):
    def __init__(self):
        super(CNN_net, self).__init__()
        # 这里是第一层
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 5, 5, padding=2, stride=1),  # 这里只弄5个滤波器
            nn.MaxPool2d(2),  # 先池化后激活。
            nn.ReLU())

        # 接下来定义第二层的卷积。
        self.layer2 = nn.Sequential(
            nn.Conv2d(5, 10, 3, padding=2, stride=1),  # 这里弄8个。
            nn.MaxPool2d(2),  # 这里是先池化，后激活。
            nn.ReLU())

        # 这里开始定义全连接层。
        self.liner1 = nn.Linear(640, 50)
        self.liner2 = nn.Linear(50, 10)  # 这里直接转化为10个，然后使用softmax函数进求解。

    def forward(self, x):  # 这里就是定义向前传播的参数了，x表示的是一个图片的特征。
        x = self.layer1(x)  # 这里进行第一步的卷积操作
        x = self.layer2(x)  # 这里把用relu函数激活后的x再一次进行卷积。
        x = x.view(-1, 8 * 8 * 10)  # 这里要将数据展开，表示的是一个一维，方便后续的全连接层。
        x = self.liner1(x)
        x = F.relu(x)
        x = self.liner2(x)
        return x


def get_it(alpha):  # 这里是获得模型的函数,传入的参数是学习率
    net = CNN_net()  # 首先还是实例化
    optimizer = torch.optim.Adam(net.parameters(), lr=alpha)  # 这里设置学习率和优化方法。
    return net, optimizer  # 这里返回的是实例化后的参数以及优化方法optimizer.


# 获取训练集
train = torchvision.datasets.FashionMNIST(root='./datasets/', train=True, download=True,
                                          transform=transforms.ToTensor())
train_data = DataLoader(dataset=train, batch_size=600, shuffle=False)
# 获取测试集
test = torchvision.datasets.FashionMNIST(root='./datasets/', train=False, download=True,
                                         transform=transforms.ToTensor())
test_data = DataLoader(dataset=test, batch_size=100, shuffle=False)

epoch = 20  # 这里表示整个神经网络迭代50次即可。
alpha = 0.01  # 这里表示的是定义学习率。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 初始化GPU,将数据放在GPU上面进行训练。
net, optimizer = get_it(alpha)  # 这里是实例化以及反向传播的优化方法
loss_function = nn.CrossEntropyLoss().to(device)  # 放在GPU上面进行训练。
loss = 0  # 初始化损失值。
net = net.to(device)  # 这里是将模型也放在GPU上面进行运算。
# （四）开始训练模型。
for i in range(epoch):  # 这里表示迭代多少次的意思。
    loss = 0
    print("正在进行第" + str(i + 1) + "次迭代")
    for a, b in train_data:
        a = a.to(device)
        b = b.to(device)  # 这里将数据全部放到GPU上面进行训练。
        optimizer.zero_grad()  # 初始化权重。
        y_pre = net(a)  # 这里得到最后的答案。
        loss_data = loss_function(y_pre, b)  # 计算出损失值。
        loss += loss_data / 600  # 这里是除以你batch_size的大小即可。
        # print("损失值为：" + str(loss_data/10000))
        loss_data.backward()  # 反向传播。
        optimizer.step()
    print("第" + str(i + 1) + "次迭代完成。" + "损失值为：" + str(float(loss)))
print("任务完成")

accurate = 0  # 初始话为0.
for a, b in test_data:
    a = a.to(device)
    b = b.to(device)  # 这里将数据全部放到GPU上面进行训练。
    finish = net(a)
    aaa, finish_pre = torch.max(finish, axis=1)
    # print(finish_pre)
    acc = (finish_pre == b)
    accurate += acc.sum()
print("测试集的准确率为：")
print(str(float(accurate/100)) + "%")